# coding=utf-8
import sys
import json
import os
from boto3.session import Session
from botocore.exceptions import ClientError
import subprocess
from reportlab.pdfgen import canvas
from reportlab.lib.units import cm
from PyPDF2 import PdfReader, PdfWriter
from reportlab.pdfbase import pdfmetrics, ttfonts

MARK_FILE = "/tmp/mark.pdf"

# support chinese
# pdfmetrics.registerFont(
#     ttfonts.TTFont(
#         "zenhei", os.path.join("/usr/share/fonts/truetype/wqy", "wqy-zenhei.ttc")
#     )
# )
# pdfmetrics.registerFont(
#     ttfonts.TTFont(
#         "microhei", os.path.join("/usr/share/fonts/truetype/wqy", "wqy-microhei.ttc")
#     )
# )

def add_watermark(pdf_file_in, pdf_file_mark, pdf_file_out):
    print(pdf_file_in, pdf_file_mark, pdf_file_out)
    pdf_output = PdfWriter()
    with open(pdf_file_in, "rb") as input_stream:
        pdf_input = PdfReader(input_stream, strict=False)

        # pageNum = pdf_input.getNumPages()
        pageNum = len(pdf_input.pages)
        pdf_watermark = PdfReader(open(pdf_file_mark, "rb"), strict=False)

        for i in range(pageNum):
            page = pdf_input.pages[i]
            page.merge_page(pdf_watermark.pages[0])
            page.compress_content_streams()
            pdf_output.add_page(page)

        with open(pdf_file_out, "wb") as f:
            pdf_output.write(f)


def _create_watermark(
    mark_text,
    pagesize=(21 * cm, 29.7 * cm),
    font="Helvetica",
    font_size=30,
    font_color=(0, 0, 0),
    rotate=0,
    opacity=1,
    density=(5 * cm, 5 * cm),
):
    c = canvas.Canvas(MARK_FILE, pagesize=pagesize)
    c.setFont(font, font_size)
    c.rotate(rotate)
    c.setStrokeColorRGB(0, 0, 0)
    r, g, b = font_color
    c.setFillColorRGB(r, g, b)
    c.setFillAlpha(opacity)

    row_gap, col_gap = density
    colN = int(pagesize[0] / col_gap)
    rowN = int(pagesize[1] / row_gap)
    x = colN * 4
    y = rowN * 4

    for i in range(y):
        for j in range(x):
            a = col_gap * (j - 2 * colN)
            b = row_gap * (i - 2 * rowN)
            c.drawString(a, b, mark_text)

    c.save()


def create_watermark(evt):
    mark_text = evt.get("mark_text", "hello ctyun cf")
    # 1cm = 28.346456692913385， defalut is A4, (21*cm, 29.7*cm)
    pagesize = evt.get("pagesize", [595.275590551181, 841.8897637795275])
    font = evt.get("font", "Helvetica")
    font_size = evt.get("font_size", 30)
    font_color = evt.get("font_color", (0, 0, 0))
    rotate = evt.get("rotate", 30)
    opacity = evt.get("opacity", 0.1)
    # default is (7*cm, 10*cm)
    density = evt.get("density", [198.4251968503937, 283.46456692913387])
    _create_watermark(
        mark_text, pagesize, font, font_size, font_color, rotate, opacity, density
    )
    print("create_watermark success!")


def handler(event, context):
    access_key = os.getenv("ext_cf_zos_ak")
    secret_key = os.getenv("ext_cf_zos_sk")
    endpoint = os.getenv("ext_cf_zos_endpoint")
    bucket = os.getenv("ext_cf_zos_bucket")
    if access_key is None:
        return "env check failed: zos ak is None"
    if secret_key is None:
        return "env check failed: zos sk is None"
    if endpoint is None:
        return "env check failed: zos endpoint is None"
    if bucket is None:
        return "env check failed: zos bucket is None"

    session = Session(access_key, secret_key)
    s3_client = session.client("s3", endpoint_url=endpoint)

    evt = json.loads(event)
    remote_pdf_file = evt["pdf_file"]
    remote_file_dir, file_name = os.path.split(remote_pdf_file)
    short_file_name, _ = os.path.splitext(file_name)

    local_save_dir = "/tmp"
    local_pdf_file = local_save_dir + "/" + file_name
    local_pdf_out_file = local_save_dir + "/" + short_file_name + "_out.pdf"

    # 下载文件
    try:
        s3_client.download_file(bucket, remote_pdf_file, local_pdf_file)
    except ClientError as e:
        return "download pdf file from zos failed: %s" % e

    # 添加水印
    create_watermark(evt)
    add_watermark(local_pdf_file, MARK_FILE, local_pdf_out_file)
    print("pdf add_watermark success!")

    # 上传文件
    try:
        s3_client.upload_file(local_pdf_out_file, bucket, short_file_name + "_out.pdf")
    except ClientError as e:
        return "upload pdf file to zos failed: %s" % e

    subprocess.check_call(["ls", "-al", "/tmp"])
    subprocess.check_call("rm -rf /tmp/*", shell=True)

    return "upload to zos success"
